import paddle
import numpy as np
import pandas as pd
from tqdm import tqdm


class FeedFroward(paddle.nn.Layer):
    def __init__(self, hidden_dim):
        super(FeedFroward, self).__init__()
        self.fc1 = paddle.nn.Linear(hidden_dim, hidden_dim // 2, bias_attr=False)
        self.fc2 = paddle.nn.Linear(hidden_dim // 2, hidden_dim, bias_attr=False)
        self.dr = paddle.nn.Dropout(0.2)

    def forward(self, feed_x):
        feed_x = self.fc1(feed_x)

        feed_x = paddle.sin(feed_x)
        feed_x = self.fc2(self.dr(feed_x))
        return feed_x


class ImageEm(paddle.nn.Layer):
    def __init__(self, k_size, hidden_dim, sp_num=3):
        super(ImageEm, self).__init__()
        self.fc1 = paddle.nn.Conv2D(in_channels=3, out_channels=sp_num * hidden_dim, bias_attr=False,
                                    kernel_size=k_size, stride=k_size)

        self.sp = sp_num

        self.feed = FeedFroward(hidden_dim)

    def forward(self, image_x):
        image_x = self.fc1(image_x)
        image_x = image_x.reshape([image_x.shape[0], self.sp, -1, 16, 16]).transpose([0, 1, 3, 4, 2])
        image_x = image_x.reshape([image_x.shape[0], -1, image_x.shape[-1]])
        image_x = self.feed(image_x)
        return image_x


class WavEm(paddle.nn.Layer):
    def __init__(self, hidden_dim, k_size, sp_num=3):
        super(WavEm, self).__init__()
        self.fc1 = paddle.nn.Conv1D(in_channels=1, out_channels=hidden_dim * sp_num, bias_attr=False,
                                    kernel_size=k_size, stride=k_size)
        self.feed = FeedFroward(hidden_dim)
        self.sp = sp_num

    def forward(self, wav_x):
        wav_x = self.fc1(wav_x)
        wav_x = wav_x.reshape([wav_x.shape[0], -1, self.sp, wav_x.shape[-1]]).transpose([0, 2, 3, 1])
        wav_x = wav_x.reshape([wav_x.shape[0], -1, wav_x.shape[-1]])
        wav_x = self.feed(wav_x)
        return wav_x


class ImageOut(paddle.nn.Layer):
    def __init__(self, hidden_dim, k_size, sp_num=3):
        super(ImageOut, self).__init__()
        self.fc1 = paddle.nn.Conv2DTranspose(in_channels=sp_num * hidden_dim, out_channels=3, bias_attr=False,
                                             kernel_size=k_size, stride=k_size)

        self.sp = sp_num

        self.feed = FeedFroward(hidden_dim)

    def forward(self, image_out_x):
        self.feed(image_out_x)
        image_out_x = image_out_x.reshape([image_out_x.shape[0], -1, 16, 16])
        image_out_x = self.fc1(image_out_x)
        return image_out_x


class WavOut(paddle.nn.Layer):
    def __init__(self, hidden_dim, k_size, sp_num=3):
        super(WavOut, self).__init__()
        self.fc1 = paddle.nn.Conv1DTranspose(in_channels=hidden_dim * sp_num, out_channels=1, bias_attr=False,
                                             kernel_size=k_size, stride=k_size)
        self.feed = FeedFroward(hidden_dim)
        self.sp = sp_num

    def forward(self, wav_out_x):
        wav_out_x = self.feed(wav_out_x)
        wav_out_x = wav_out_x.reshape([wav_out_x.shape[0], self.sp, -1, wav_out_x.shape[-1]])
        wav_out_x = wav_out_x.transpose([0, 1, 3, 2])
        wav_out_x = wav_out_x.reshape([wav_out_x.shape[0], -1, wav_out_x.shape[-1]])
        wav_out_x = self.fc1(wav_out_x)
        return wav_out_x


class CVFroward(paddle.nn.Layer):
    def __init__(self, hidden_dim, k_size):
        super(CVFroward, self).__init__()
        self.fc1 = paddle.nn.Conv1D(in_channels=hidden_dim, out_channels=hidden_dim // 2, bias_attr=False,
                                    kernel_size=k_size, stride=1)
        self.fc2 = paddle.nn.Conv1D(in_channels=hidden_dim // 2, out_channels=hidden_dim, bias_attr=False,
                                    kernel_size=k_size, stride=1)

        self.pad = k_size - 1
        self.dr = paddle.nn.Dropout(0.2)

    def forward(self, feed_x):
        feed_x = feed_x.transpose([0, 2, 1])
        zero_pad = paddle.zeros(shape=[feed_x.shape[0], feed_x.shape[1], self.pad], dtype=feed_x.dtype)
        feed_x = paddle.concat([zero_pad, feed_x], -1)
        feed_x = self.fc1(feed_x)

        feed_x = paddle.sin(feed_x)
        zero_pad = paddle.zeros(shape=[feed_x.shape[0], feed_x.shape[1], self.pad], dtype=feed_x.dtype)

        feed_x = paddle.concat([zero_pad, feed_x], -1)
        feed_x = self.fc2(self.dr(feed_x))
        feed_x = feed_x.transpose([0, 2, 1])

        return feed_x


class CvfBlock(paddle.nn.Layer):
    def __init__(self, hidden_dim, k_size):
        super(CvfBlock, self).__init__()
        self.cvf = CVFroward(hidden_dim, k_size)
        self.feed = FeedFroward(hidden_dim)

    def forward(self, ll):
        ll += self.cvf(ll)
        ll += self.feed(ll)
        return ll


class CVL(paddle.nn.Layer):
    def __init__(self, hidden_dim, k_max):
        super(CVL, self).__init__()
        self.cvf = paddle.nn.LayerList([CvfBlock(hidden_dim, one_ksize) for one_ksize in range(3, k_max)])
        # self.cvl_out = paddle.nn.Linear(len(list(range(3,k_max))), 1,bias_attr=False)
        self.layer_nor = paddle.nn.LayerNorm(hidden_dim)
        self.group_nor = paddle.nn.GroupNorm(num_channels=k_max - 3, num_groups=k_max - 3)

    def forward(self, x):
        x = paddle.concat([cvf(x).unsqueeze(-1) for cvf in self.cvf], -1)
        # x = self.cvl_out(x).squeeze(-1)
        x = self.group_nor(x.transpose([0, 3, 1, 2]))
        x = self.layer_nor(paddle.sum(x, 1))
        return x


class CVFLora(paddle.nn.Layer):
    def __init__(self, hidden_dim, layers=2):
        super(CVFLora, self).__init__()
        self.lora = paddle.nn.LayerList([FeedFroward(hidden_dim) for _ in range(layers)])
        self.layer_nor = paddle.nn.LayerNorm(hidden_dim)
        self.group_nor = paddle.nn.GroupNorm(num_channels=layers, num_groups=layers)

    def forward(self, x):
        x = paddle.concat([lora(x).unsqueeze(-1) for lora in self.lora], -1)
        x = self.group_nor(x.transpose([0, 3, 1, 2]))
        x = self.layer_nor(paddle.sum(x, 1))
        return x


class CVFBlock(paddle.nn.Layer):
    def __init__(self, hidden_dim, k_size_max):
        super(CVFBlock, self).__init__()

        self.cvl = CVL(hidden_dim * 3, k_size_max)

        self.lora = CVFLora(hidden_dim, 5)
        self.lora1 = CVFLora(hidden_dim, 5)

        self.layer_nor = paddle.nn.LayerNorm(hidden_dim)

    def forward(self, x, state=0):
        x = paddle.concat([x, x, x], -1)

        x = self.cvl(x)
        x = x.reshape([x.shape[0], x.shape[1], 3, -1])
        c = x[:, :, 0]
        v = x[:, :, 1]
        f = x[:, :, 2] + state

        vf = self.lora(self.layer_nor(v + f))
        cvf = self.lora1(self.layer_nor(c + vf))

        return cvf


class CVF(paddle.nn.Layer):
    def __init__(self, voc_size, hidden_dim, k_size_max, multi_layers):
        super(CVF, self).__init__()

        self.emx = paddle.nn.Embedding(voc_size, hidden_dim)
        self.cvf_block = paddle.nn.LayerList([CVFBlock(hidden_dim, k_size_max) for _ in range(multi_layers)])

        self.out_layer = paddle.nn.Linear(hidden_dim, voc_size, bias_attr=False)

    def forward(self, x):
        x = self.emx(x)
        cvf_state = 0
        for i, cvf in enumerate(self.cvf_block):
            cvf_state = cvf(x, cvf_state)

        x = self.out_layer(cvf_state)

        return x


class CvfSpToImage(paddle.nn.Layer):
    def __init__(self, voc_size, hidden_dim, k_size_max, multi_layers):
        super(CvfSpToImage, self).__init__()

        self.emx = paddle.nn.Embedding(voc_size, hidden_dim)
        self.cvf_block = paddle.nn.LayerList([CVFBlock(hidden_dim, k_size_max) for _ in range(multi_layers)])

        self.out_layer = ImageOut(hidden_dim, 16)

    def forward(self, x):
        x = self.emx(x)
        cvf_state = 0
        for i, cvf in enumerate(self.cvf_block):
            cvf_state = cvf(x, cvf_state)

        x = self.out_layer(cvf_state)

        return x


class CvfSpToWav(paddle.nn.Layer):
    def __init__(self, voc_size, hidden_dim, k_size_max, multi_layers):
        super(CvfSpToWav, self).__init__()

        self.emx = paddle.nn.Embedding(voc_size, hidden_dim)
        self.cvf_block = paddle.nn.LayerList([CVFBlock(hidden_dim, k_size_max) for _ in range(multi_layers)])

        self.out_layer = WavOut(hidden_dim, 320)

    def forward(self, x):
        x = self.emx(x)
        cvf_state = 0
        for i, cvf in enumerate(self.cvf_block):
            cvf_state = cvf(x, cvf_state)

        x = self.out_layer(cvf_state)

        return x


class CvfImageToSp(paddle.nn.Layer):
    def __init__(self, voc_size, hidden_dim, k_size_max, multi_layers):
        super(CvfImageToSp, self).__init__()

        self.emx = ImageEm(16, hidden_dim)
        self.cvf_block = paddle.nn.LayerList([CVFBlock(hidden_dim, k_size_max) for _ in range(multi_layers)])

        self.out_layer = paddle.nn.Linear(hidden_dim, voc_size, bias_attr=False)

    def forward(self, x):
        x = self.emx(x)
        cvf_state = 0
        for i, cvf in enumerate(self.cvf_block):
            cvf_state = cvf(x, cvf_state)

        x = self.out_layer(cvf_state)

        return x


class CvfWavToSp(paddle.nn.Layer):
    def __init__(self, voc_size, hidden_dim, k_size_max, multi_layers):
        super(CvfWavToSp, self).__init__()

        self.emx = WavEm(320, hidden_dim)
        self.cvf_block = paddle.nn.LayerList([CVFBlock(hidden_dim, k_size_max) for _ in range(multi_layers)])

        self.out_layer = paddle.nn.Linear(hidden_dim, voc_size, bias_attr=False)

    def forward(self, x):
        x = self.emx(x)
        cvf_state = 0
        for i, cvf in enumerate(self.cvf_block):
            cvf_state = cvf(x, cvf_state)

        x = self.out_layer(cvf_state)

        return x


def sequence_clarity_paddle(x, short=2, long=10):
    """

    :param x:
    :param short:
    :param long:
    :return:
    """
    score_sort = paddle.sort(paddle.nn.functional.softmax(x, -1), -1)
    basic_score = short ** (paddle.arange(long)) / sum(short ** paddle.arange(long))
    basic_score = paddle.concat([paddle.zeros([score_sort.shape[-1] - long]), basic_score])
    basic_score = paddle.zeros([score_sort.shape[0], score_sort.shape[1], 1]) + basic_score
    ming_que_du = paddle.mean(paddle.nn.functional.cosine_similarity(score_sort, basic_score, -1))
    return ming_que_du


class SequenceClarityLoss(paddle.nn.Layer):
    def __init__(self):
        super(SequenceClarityLoss, self).__init__()
        self.lss = paddle.nn.CrossEntropyLoss()

    def forward(self, x, loss1):
        score = sequence_clarity_paddle(x)
        ss = self.lss(x, loss1) * (1 - score)
        return ss, score


def em_heading_train_and_sample():
    print("*" * 100)

    # net = ReNetAll(12935, 256, 8)
    net = CVF(1 + 256 * 12, 128, 8, 9)

    x = paddle.to_tensor([
        np.random.randint(1, 1000, 5120),
        np.random.randint(1, 1000, 5120),
    ], dtype='int64')

    # 模拟训练

    # loss_f = paddle.nn.CrossEntropyLoss()
    loss_f = SequenceClarityLoss()

    opt = paddle.optimizer.Adam(parameters=net.parameters(), learning_rate=0.0003)
    bar = tqdm(range(60))
    loss_list = []
    for epoch in bar:

        for i in range(0, x.shape[-1], 256):
            j = i + 256

            out1 = net(x[:, i + 1:j - 1])

            loss, score = loss_f(out1, x[:, i + 2:j])

            loss_list.append(loss.item())

            bar.set_description("epoch---{}--loss--{:.5f}---mqd---{:.5f}".format(epoch, loss.item(), score.item()))

            loss.backward()
        opt.step()
        opt.clear_grad()
    pd.to_pickle({"data": loss_list}, "loss_list.pandas_pickle")


def run_sp_all():
    net = ImageEm(16, 128)
    image_out = net(paddle.rand([1, 3, 256, 256]))

    net = WavEm(128, 320)
    wav_out = net(paddle.rand([1, 1, 160000]))

    net = ImageOut(128, 16)
    net(image_out)

    net = WavOut(128, 320)
    net(wav_out)


def model_param_count():
    model = CvfSpToWav(10000,512, 8, 46)
    sum(paddle.numel(p) for p in model.parameters())
    print("CvfSpToWav:{:.3f}B".format(sum(paddle.numel(p) for p in model.parameters()).item() / 1000000000))

    model = CvfSpToImage(65535, 128, 8, 8)
    sum(paddle.numel(p) for p in model.parameters())
    print("CvfSpToImage:{:.3f}B".format(sum(paddle.numel(p) for p in model.parameters()).item() / 1000000000))

    model = CvfImageToSp(65535, 128, 8, 8)
    sum(paddle.numel(p) for p in model.parameters())
    print("CvfImageToSp:{:.3f}B".format(sum(paddle.numel(p) for p in model.parameters()).item() / 1000000000))

    model = CvfWavToSp(65535, 128, 8, 8)
    sum(paddle.numel(p) for p in model.parameters())
    print("CvfWavToSp:{:.3f}B".format(sum(paddle.numel(p) for p in model.parameters()).item() / 1000000000))


# 进行模型训练和预测
if __name__ == '__main__':
    model_param_count()
    # run_sp_all()
    # em_heading_train_and_sample()
